import numpy as np 
from matplotlib import pyplot as plt
from mpl_toolkits.mplot3d import Axes3D

## n 样本数
## k 聚类中心数
## dimension 样本维度
## file_name 数据集
def K_means(n,k,dimension,file_name):
    fp=open(file_name,'r')
    Lid=[0 for i in range(n)]
    L=[]
    lines=fp.readlines()
    Lc=['r','b','g','y','m']
    #### 读入数据
    for line in lines:
        row=list(map(float,line.strip('\n').split()))
        L.append(row)
    Lk=[[0.0 for i in range(c)] for i in range(k)]
    #### 确定初始聚类中心、默认数据前k个
    for i in range(k):
        for j in range(dimension):
            Lk[i][j]=L[i][j]
    #### 循环确定聚类中心
    while(True):
        for i in range(n):
            tmpans=1000000.0
            for j in range(k):
                tmp=0.0
                for q in range(dimension):
                    tmp+=(L[i][q]-Lk[j][q])*(L[i][q]-Lk[j][q])
                if tmp<tmpans:
                    tmpans=tmp
                    Lid[i]=j
        Lnum=[0 for i in range(k)]
        Lnew=[[0.0 for i in range(dimension)] for i in range(k)]
        for i in range(n):
            Lnum[Lid[i]]+=1
            for j in range(dimension):
                Lnew[Lid[i]][j]+=L[i][j]
        flag=0
        for i in range(k):
            for j in range(dimension):
                tmp=Lnew[i][j]/Lnum[i]
                if Lk[i][j]!=tmp:
                    flag+=1
                    Lk[i][j]=tmp
        if flag==0: break

    #### 利用相关库将数据可视化
    ## 维度为2
    for i in range(n):
        plt.scatter(L[i][0],L[i][1],10,Lc[Lid[i]])
    plt.savefig("2D.jpg")
    plt.draw()
    plt.pause(10)    
    ## 维度为3
    ax = plt.subplot(projection='3d') 
    for i in range(n):
        ax.scatter(L[i][0],L[i][1],L[i][2],s=5,c=Lc[Lid[i]])
    plt.savefig("3D.jpg")
    plt.show()
    plt.pause(5)

